# Copyright © 2022 BAAI. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License")

import os
import shlex
import socket
import argparse
import subprocess
from collections import OrderedDict

import torch


def inspect_dict(input, level=0):
    """ Print the keys and types of a dictionary recursively.

    Args:
        dict: dictionary to be inspected
    """
    for key in input:
        if isinstance(input[key], dict):
            print("  " * level, "Key {} of type {}".format(key, type(input[key])), flush=True)
            inspect_dict(input[key], level + 1)
        elif isinstance(input[key], torch.Tensor):
            print("  " * level, "Key {} of shape {}".format(key, input[key].shape), flush=True)
        else:
            print("  " * level, "Key {} of type {}".format(key, type(input[key])), flush=True)


def compare_dicts(first_dict, second_dict):
    """ Compare two dictionaries recursively.

    Args:
        frist_dict: first dictionary
        second_dict: second dictionary
    
    Returns:
        True if the dictionaries are equal, False otherwise
    """
    print("Comparing {} and {}".format(first_dict.keys(), second_dict.keys()))
    second_keys = set(second_dict.keys())
    second_visited_keys = set()
    for key in first_dict:
        if key not in second_dict:
            print("Key {} not found in second dict".format(key))
            return False
        elif type(first_dict[key]) == type(second_dict[key]):
            second_visited_keys.add(key)
            if isinstance(first_dict[key], dict):
                if not compare_dicts(first_dict[key], second_dict[key]):
                    return False
            else:
                if not torch.equal(first_dict[key], second_dict[key]):
                    print("Key {} not equal".format(key))
                    return False  
        else:
            print("Key {} not equal".format(key))
            return False
    if len(second_keys) != len(second_visited_keys):
        print("Second dict contains extra keys {}".format(second_keys - second_visited_keys))
        return False
    print("Dictionaries {} and {} are equal".format(first_dict.keys(), second_dict.keys()))
    return True


def get_args():
    """Parse commandline."""

    parser = argparse.ArgumentParser(description='Pull remote checkpoints')
    parser.add_argument('--ckpts', nargs='+', required=True,
                        help='Path to model checkpoints to be compared')
    parser.add_argument('--action', default='compare',
                       choices=['compare', 'inspect'],
                       help='')
    args = parser.parse_args()

    return args
                        

if __name__ == '__main__':
    args = get_args()

    if args.action == 'compare':
        first_ckpt = torch.load(args.ckpts[0], map_location='cpu')
        for ckpt_path in args.ckpts[1:]:
            second_ckpt = torch.load(ckpt_path, map_location='cpu')
            if not compare_dicts(first_ckpt['model'], second_ckpt['model']):
                print("Checkpoints {} and {} are not equal".format(args.ckpts[0], ckpt_path))
                exit(1)
        print("Checkpoints are equal")
    elif args.action == 'inspect':
        for ckpt_path in args.ckpts:
            ckpt = torch.load(ckpt_path, map_location='cpu')
            inspect_dict(ckpt)